package com.swt.znwd;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.model.perceptron.*;
import com.hankcs.hanlp.model.perceptron.tagset.NERTagSet;
import com.hankcs.hanlp.model.perceptron.tagset.POSTagSet;
import com.hankcs.hanlp.model.perceptron.tagset.TagSet;
import com.hankcs.hanlp.seg.Segment;
import com.hankcs.hanlp.seg.common.Term;
import com.hankcs.hanlp.utility.Predefine;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author BleethNie
 * @version 1.0
 * @date 2020-07-24 14:43
 **/

@Slf4j
public class NerAUtil {


    private static PerceptronSegmenter cwsSegment = null;


    private static PerceptronPOSTagger posTagger = null;


    private static PerceptronNERecognizer neRecognizer = null;


    /*
     *模型加载
     */
    public static String startCorpusPath = System.getProperty("user.dir") + "/doc/corpus";
    public static String startModelPath = System.getProperty("user.dir") + "/doc/model";

    public static String cwsTrainingFile = startCorpusPath + "/corpus_train.txt";
    public static String cwsTestFile = startCorpusPath + "/corpus_test.txt";

    public static String cwsSaveFile = startModelPath + "/cws.bin";
    public static String posSaveFile = startModelPath + "/pos.bin";
    public static String nerSaveFile = startModelPath + "/ner.bin";


    //加载标注集
    public static String tagFile_time = System.getProperty("user.dir")+"/corpus/time_001.txt";

    public static String tagFile_address = System.getProperty("user.dir")+"/corpus/address_001.txt";


    /**
     * 初始化使用模型
     *
     * @throws IOException
     */
    private static void init() throws IOException {

        cwsSegment = new PerceptronSegmenter(cwsSaveFile);

        posTagger = new PerceptronPOSTagger(posSaveFile);

        neRecognizer = new PerceptronNERecognizer(nerSaveFile);
        log.info("加载实体提取模型成功...");
    }


    /**
     * 训练模型
     *
     * @throws IOException
     */
    public static void trainModel() throws IOException {
        PerceptronTrainer cwsTrainer = new CWSTrainer();
        PerceptronTrainer.Result train_result = cwsTrainer.train(cwsTrainingFile, cwsTestFile, cwsSaveFile);
        System.out.println(train_result.getAccuracy());

        PerceptronTrainer postrainer = new POSTrainer() {
            @Override
            protected TagSet createTagSet() {
                POSTagSet tagSet = new POSTagSet();
                tagSet.add("time");
                tagSet.add("address");
                return tagSet;
            }
        };
        PerceptronTrainer.Result posmodel = postrainer.train(cwsTrainingFile, cwsTestFile, posSaveFile, 0.1, 5, 2);
        System.out.println("模型训练时的准确率为:" + posmodel.getAccuracy());


        PerceptronTrainer nertrainer = new NERTrainer() {
            @Override
            protected TagSet createTagSet() {
                NERTagSet tagSet = new NERTagSet();
                tagSet.nerLabels.add("time");
                tagSet.nerLabels.add("address");
                return tagSet;
            }
        };
        PerceptronTrainer.Result nerModel = nertrainer.train(cwsTrainingFile, cwsTestFile, nerSaveFile, 0.1, 4, 4);
    }

    /**
     * 加载语料
     *
     * @throws SQLException
     */
    public static void loadTrainCorpus() throws SQLException {
        List<Entity> entityList = Db.use().query("select * from t_xab_lb where   zz_corpus_type = '1' ");
        List<String> list = new ArrayList<>();
        for (int i = 0; i < entityList.size(); i++) {
            Entity entity = entityList.get(i);
            String zz_corpus = entity.getStr("zz_corpus");
            list.add(zz_corpus);
        }
        FileUtil.writeLines(list, cwsTrainingFile, "UTF-8", false);
    }

    /**
     * 加载语料
     *
     * @throws SQLException
     */
    public static void loadTestCorpus() throws SQLException {
        List<Entity> entityList = Db.use().query("select * from t_xab_lb where   zz_corpus_type = '3' ");
        List<String> list = new ArrayList<>();
        for (int i = 0; i < entityList.size(); i++) {
            Entity entity = entityList.get(i);
            String zz_corpus = entity.getStr("zz_corpus");
            list.add(zz_corpus);
        }
        FileUtil.writeLines(list, cwsTestFile, "UTF-8", false);
    }


    /**
     * 预处理
     *
     * @throws SQLException
     */
    public static void preHandler() throws SQLException {
        Predefine.HANLP_PROPERTIES_PATH = System.getProperty("user.dir") + "/config/hanlp.properties";


        Segment segment = HanLP.newSegment("perceptron");

        List<Entity> entityList = Db.use().query("select * from t_xab_lb where zz_corpus_type = '1' ");

        for (int i = 0; i < entityList.size(); i++) {
            Entity entity = entityList.get(i);
            String zyaq = entity.getStr("zyaq");
            if (StrUtil.isEmpty(zyaq)) {
                continue;
            }
            List<Term> seg = segment.seg(zyaq);
            String segContent = Convert.toStr(seg);
            segContent = StrUtil.removePrefix(segContent, "[");
            segContent = StrUtil.removeSuffix(segContent, "]");
            segContent = StrUtil.replace(segContent, ", ", " ");


            Db.use().update(
                    Entity.create().set("zz_corpus", segContent), //修改的数据
                    Entity.create("t_xab_lb").set("zz_id", entity.getInt("zz_id")) //where条件
            );
        }

    }


    /**
     * 学习词条
     *
     * @param text
     */
    public String learn(String text) {
        List<String> segment = cwsSegment.segment(text);
        String[] posTags = posTagger.tag(ArrayUtil.toArray(segment, String.class));
        String[] nerTags = neRecognizer.recognize(ArrayUtil.toArray(segment, String.class), posTags);
        return "";
    }


    public static void test() throws IOException, SQLException {
        init();
        List<Entity> entityList = Db.use().query("select * from t_xab_lb  ORDER BY RAND() limit 10     ");
        for (int i = 0; i < entityList.size(); i++) {
            Entity entity = entityList.get(i);
            String zyaq = entity.getStr("zyaq");
            if (StrUtil.isEmpty(zyaq)) {
                continue;
            }
            List<String> segmentList = cwsSegment.segment(zyaq);

            String[] tag = posTagger.tag(ArrayUtil.toArray(segmentList, String.class));
            List<String> tagList = Convert.toList(String.class, tag);

            String[] ner = neRecognizer.recognize(ArrayUtil.toArray(segmentList, String.class), tag);
            List<String> nerList = Convert.toList(String.class, ner);

            List<String> sList = new ArrayList<>();
            for (int segIndex = 0; segIndex < segmentList.size(); segIndex++) {
                String s = segmentList.get(segIndex);
                String t = tag[segIndex];
                String n = ner[segIndex];
                sList.add( s + "/" + t + "-" + n);
            }

//            System.err.println(sList);
            System.err.println(segmentList);
            extract("nr", segmentList, tagList, nerList);
            extract("time", segmentList, tagList, nerList);
            extract("address", segmentList, tagList, nerList);

            System.err.println("========================================================" +
                    "===================================================================");
        }
    }


    public static void load(String tagFile) {
        List<String> tagLine = FileUtil.readLines(tagFile, "UTF-8");
       // System.out.println("大小是"+tagLine.size());
        //80%的是训练集，20%测试集
        List<String> trainLine = new ArrayList<>();
        List<String> testLine = new ArrayList<>();
        for (int i = 0; i < tagLine.size(); i++) {
            if (i % 8 == 0) {
                String line = tagLine.get(i);
                testLine.add(line);
            } else {
                String line = tagLine.get(i);
                trainLine.add(line);
            }
        }

        FileUtil.writeLines(testLine, cwsTestFile, "UTF-8", true);
        FileUtil.writeLines(trainLine, cwsTrainingFile, "UTF-8", true);

    }


    public static void main(String[] args) throws SQLException, IOException {
        FileUtil.del(cwsTestFile);
        FileUtil.del(cwsTrainingFile);

        load(tagFile_time);
        load(tagFile_address);

        trainModel();
        test();


    }

    //address、nr、time
    public static void extract(String tag, List<String> segList, List<String> posList, List<String> nerList) {
        if ("nr".equals(tag)) {
            List<String> nrList = new ArrayList<>();
            for (int i = 0; i < posList.size(); i++) {
                String seg = segList.get(i);
                String pos = posList.get(i);
                if (pos.equals("nr")) {
                    nrList.add(seg);
                }
            }
            System.err.println("人名实体:" + nrList);
        }
        if ("time".equals(tag)) {
            List<String> timeList = new ArrayList<>();
            String timeSeg = "";
            for (int i = 0; i < posList.size(); i++)    {
                String seg = segList.get(i);
                //System.out.println("......"+seg);
                String pos = posList.get(i);
                String ner = nerList.get(i);
                if (ner.endsWith("time")) {
                    timeSeg = timeSeg + seg;
                    //System.out.println("===="+timeSeg);
                }
                //E-time 表示以time 结尾
                if (ner.equals("E-time")) {
                    timeList.add(timeSeg);
                    timeSeg = "";
                }
            }
            System.err.println("时间实体:" + timeList);
        }
        if ("address".equals(tag)) {
            List<String> addressList = new ArrayList<>();
            String addressSeg = "";
            for (int i = 0; i < posList.size(); i++) {
                String seg = segList.get(i);
                String pos = posList.get(i);
                String ner = nerList.get(i);
                if (ner.endsWith("address")) {
                    addressSeg = addressSeg + seg;
                }
                if (ner.equals("E-address")) {
                    addressList.add(addressSeg);
                    addressSeg = "";
                }
            }
            System.err.println("地址实体:" + addressList);
        }
    }


    //作案手段提取
    public static void zasd(String content) {
          List<String> zList=new ArrayList<>();


    }

    //报案人特征提取
    public static void bar(String content) {
        //姓名

        //性别

        //住址

        //年龄

        //手机号

    }

    //受害人特征提取
    public static void shr(String content) {
        //姓名

        //性别

        //住址

        //年龄

        //手机号

    }


    //嫌疑人特征提取
    public static void xyr(String content) {
        //姓名

        //性别

        //住址

        //年龄

        //手机号


    }


    //损失金额特征提取
    public static void ssje(String content) {



    }

    //作案工具特征提取
    public static void zagj(String content) {

    }



}
